import numpy as np
import copy

from matplotlib import pyplot as plt

map = {
    0: [0, 0], 1: [0, 1], 2: [0, 2], 3: [0, 3], 4: [1, 0], 5: [1, 2], 6: [1, 3], 7: [2, 0], 8: [2, 1], 9: [2, 2],
    10: [2, 3]
}


def setting(x, a):
    step = back = 0
    if x < 5 and x >= 2:
        step += 1
    if x >= 5 and x < 8:
        back += 1

    res = {}  # 返回结果是一个字典，key是下一个状态，value是下一个状态的概率

    # 分别表示向上，左，右，下的概率，初始为0，1
    forward = left = right = backward = 0.1
    # 根据策略a来分别赋值，比如a==0时
    # 向上的概率就是0.8，向下的概率就是0
    if a == 0:
        forward = 0.8
        backward = 0
    elif a == 1:
        left = 0.8
        right = 0
    elif a == 2:
        right = 0.8
        left = 0
    else:
        forward = 0
        backward = 0.8

    # 接下来就是先获取map
    # 然后tmp就是状态x的实际坐标
    Dict = copy.deepcopy(map)
    tmp = Dict[x]
    tmp[0] += 1  # 表示上面一格的坐标
    # 根据题目给的地图可以看出，上面一格的状态就是当前状态加4，同时为了保证不超过界限所以要%11
    # 如果“上面一格的状态”的坐标等于tmp就表示上面那一格就是可达的，反之不可达，即把概率置为0
    # step是因为坐标中有一个障碍物，导致“上面一格的状态就是当前状态加4”不绝对成立，所以作特判处理
    if tmp == map[(x + 4 - step) % 11]:
        res[(x + 4 - step) % 11] = forward
    else:
        res[(x + 4 - step) % 11] = 0

    Dict = copy.deepcopy(map)
    tmp = Dict[x]
    tmp[1] -= 1  # 表示向左移动一格
    if tmp == map[(x - 1 + 11) % 11]:
        res[(x - 1 + 11) % 11] = left
    else:
        res[(x - 1 + 11) % 11] = 0

    Dict = copy.deepcopy(map)
    tmp = Dict[x]
    tmp[1] += 1
    if tmp == map[(x + 1) % 11]:
        res[(x + 1) % 11] = right
    else:
        res[(x + 1) % 11] = 0

    Dict = copy.deepcopy(map)
    tmp = Dict[x]
    tmp[0] -= 1
    if tmp == map[(x - 4 + 11 + back) % 11]:
        res[(x - 4 + 11 + back) % 11] = backward
    else:
        res[(x - 4 + 11 + back) % 11] = 0

    res[x] = 1 - (res[(x + 4 - step) % 11] + res[(x - 1 + 11) % 11] + res[(x + 1) % 11] + res[(x - 4 + 11 + back) % 11])

    return res


# 归一化处理
def normal(P):
    return np.abs(P) / np.sum(np.abs(P))


def X2RowCol(x):
    return map[x]


class Env:

    def __init__(self):
        self.N = 11  # 表示状态数量，即题目中的11个格子
        self.A = np.arange(4)  # 0：向上，1：向左，2：向右，3：向下
        self.X = np.arange(self.N)  # 长为N的数组
        self.P = None  # 转移矩阵
        self.R = None  # 报酬向量，即长为N的数组，比如R[0]就表示状态0位置可获得的报酬（报酬是啥？开题目）
        self.makeP()  # 定义转义矩阵
        self.makeR()  # 定义报酬向量
        self.Gamma = 1  # 折扣，折扣为1所以报酬不打折
        self.StartState = 0  # 定义小车的初始状态
        self.EndStates = [6, 10]  # 结束状态，当小车移动到该位置是表示结束

    def action(self, x, a):
        # 环境模型通过action函数告知Agent报酬以及新状态
        x_ = np.random.choice(self.N, p=self.P[x, a, :])  # p是概率向量，根据p中的概率随机从[0,N-1]中取一个值
        return x_

    def makeP(self):
        # 初始化转义矩阵 P，这里假设 P[x, a, x_] 表示在状态 x 下执行动作 a 后转移到状态 x_ 的概率
        self.P = np.zeros((self.N, self.A.size, self.N))

        # 根据环境逻辑填充 P 矩阵
        for x in range(self.N):
            for a in range(self.A.size):
                res = setting(x, a)
                for key in res:
                    self.P[x, a, key] = res[key]

    def makeR(self):
        # 根据题目规定直接的到报酬向量
        self.R = [-0.04, -0.04, -0.04, -0.04, -0.04, -0.04, -1, -0.04, -0.04, -0.04, 1]


# 被动学习-时序差分方法TD Learning
class TD:
    def __init__(self, E):
        self.E = E  # Env（环境模型）
        self.Alpha = 0.5  # α是一个权重，称为学习率。
        # 被动学习就是根据一个给的的策略，让机器人学习状态的价值函数的,这个Pi就是给的策略比如Pi[0]表示在状态0时选择的策略
        self.Pi = [0, 1, 1, 3, 0, 0, 0, 2, 2, 2, 0]
        # 这个U就是最后求得的价值函数，即U[i]就表示状态i的价值（这个价值可以理解为权重）
        self.U = [0, 0, 0, 0, 0, 0, -1, 0, 0, 0, 1]

    def train(self):
        x = np.random.choice([0, 1, 2, 3, 4, 5, 7, 8, 9])  # 随机获取一个状态，除了6和10
        while x not in self.E.EndStates:
            a = self.Pi[x]  # 获取走下一步的策略
            _x = self.E.action(x, a)  # x 根据策略 a 可能到达的状态
            r = self.E.R[x]  # 获取状态x的报酬
            # 这个就是时序差分迭代的公式
            self.U[x] = self.U[x] + self.Alpha * (r + self.E.Gamma * self.U[_x] - self.U[x])
            x = _x

    def plot_value_function(self):
        plt.plot(self.U)
        plt.xlabel('State')
        plt.ylabel('Value')
        plt.title('Value Function')
        plt.show()


# 主动学习-Q Learning
class QLearning():
    def __init__(self, E):
        self.E = E  # Env（环境模型）
        self.Alpha = 0.5  # α是一个权重，称为学习率。
        self.Q = np.ones((11, 4)) / 4  # 一个11X4的矩阵，初始时为1，这里除以了4，所以初始都是0.25
        self.Q[10, :] = 1  # 结束状态定死了，表示到该点状态就是结束
        self.Q[6, :] = -1
        self.Pi = [0, 1, 1, 3, 0, 0, 0, 2, 2, 2, 0]

    def train(self):
        x = np.random.choice([0, 1, 2, 3, 4, 5, 7, 8, 9])
        while x not in self.E.EndStates:
            # normal函数是用来归一化的
            P = normal(self.Q[x])  # 由行为价值函数Q，得到状态x的行为价值向量（长度为4，代表每个行为的权重）
            a = np.random.choice(4, p=P)  # 根据P中的权重，随机选取一个数作为策略（0：向上，1：向左，2：向右，3：向下）
            _x = self.E.action(x, a)  # x 根据策略 a 可能到达的状态
            r = self.E.R[x]  # x 状态下得到的报酬
            # 这个是时序差分迭代的公式
            self.Q[x, a] = self.Q[x, a] + self.Alpha * (r + self.E.Gamma * np.max(self.Q[_x]) - self.Q[x, a])
            x = _x  # 进入到下一个状态

    def GetPi(self):
        self.Pi = np.argmax(self.Q, axis=1)
        return self.Pi


# 价值函数的线性逼近
class FTD():
    def __init__(self, E):
        self.w = np.array([0.5, 0.5, 0.5])
        self.E = E
        self.Alpha = 0.001
        self.Pi = [0, 1, 1, 3, 0, 0, 0, 2, 2, 2, 0]
        self.SumU = [0, 0, 0, 0, 0, 0, -1, 0, 0, 0, 1]

    def U(self, x):
        if x == 10:
            return 1
        if x == 6:
            return -1
        (row, col) = X2RowCol(x)  # 将状态 x 映射到相应的行列坐标,上面定义的map就包含各个点的实际坐标
        # 返回向量的点积结果
        # 线性逼近方法,通过调整权重w来拟合真实的价值函数
        return np.dot(np.array([1, row, col]), self.w)

    # 状态 x 对应的价值函数梯度。用于梯度下降法
    def dU(self, x):
        (row, col) = X2RowCol(x)
        return np.array([1, row, col])

    def train(self):
        x0 = np.random.choice([0, 1, 2, 3, 4, 5, 7, 8, 9])  # 随机选择一个状态
        a0 = self.Pi[x0]  # 获取走下一步的策略
        Rsum = self.E.R[x0]  # 初始化回报为初始状态 x0 的即时奖励
        # 初始化状态和动作为初始状态和初始动作
        x = x0
        a = a0
        gamma = self.E.Gamma
        while x not in self.E.EndStates:
            x = self.E.action(x, a)  # x 根据策略 a 可能到达的状态
            Rsum += gamma * self.E.R[x]  # 更新回报，加上新状态 x 的即时奖励乘以折扣因子
            a = self.Pi[x]
            gamma *= self.E.Gamma
        # 接下来这一步就运用了梯度下降法
        # self.dU(x0): 状态 x0 对应的价值函数梯度。
        self.w = self.w + self.Alpha * (Rsum - self.U(x0)) * self.dU(x0)

    def plot_value_function(self):
        plt.plot(self.SumU)
        plt.xlabel('State')
        plt.ylabel('Value')
        plt.title('Value Function')
        plt.show()


if __name__ == '__main__':
    e = Env()
    flag = cnt = 10000
    # while flag == 1:
    #     print("输入T是被动学习价值函数，输入F是线性逼近价值函数")
    #     select = input("请输入：")
    #     cnt = int(input("训练次数："))
    #
    #     if select == "T":
    #         td = TD(E=e)
    #         for i in range(cnt):
    #             td.train()
    #         # 价值函数的定义是从状态x出发，按照策略连续移动所获得的报酬总和的期望值
    #         print("被动学习价值函数:")
    #         for i in range(11):
    #             print(td.U[i])
    #
    #         td.plot_value_function()
    #         flag = 0
    #     elif select == "F":
    #         ftd = FTD(E=e)
    #         # 价值函数的定义是从状态x出发，按照策略连续移动所获得的报酬总和的期望值
    #         print("线性逼近价值函数:")
    #         for i in range(cnt):
    #             ftd.train()
    #         print("U:")
    #         for i in range(11):
    #             ftd.SumU[i]=ftd.U(i)
    #             print(ftd.SumU[i])
    #         ftd.plot_value_function()
    #         flag = 0
    #     else:
    #         print("请重新输入！")

    ql = QLearning(E=e)
    for i in range(cnt):
        ql.train()
    print("行为价值函数Q:", ql.Q)

    print("按照行为价值函数每一个状态取最高的价值得到的策略:", ql.GetPi())
